// Copyright 2020 The Tilt Brush Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

using System.Collections.Generic;
using UnityEngine;

namespace TiltBrush
{

    class FlatGeometryBrush : GeometryBrush
    {
        const int kVertsInSolid = 4;
        const int kSharedVertsInSolidPair = 2;
        const int kTrisInSolid = 2;
        const int kSharedTrisInSolidPair = 0;
        const int kMinimumKnotsAfterBreak = 6; // shortest stroke allowed after a break, in knots.

        protected const int BR = 0; // back right  (top)
        protected const int BL = 1; // back left   (top)
        protected const int FR = 2; // front right (top)
        protected const int FL = 3; // front left  (top)

        protected enum UVStyle
        {
            Distance,
            Stretch
        };

        protected List<float> m_sizes;

        const float kSolidMinLengthMeters_PS = 0.002f;
        const float kMinMoveLengthMeters_PS = 5e-4f;
        const float kBreakAngleScalar = 2.0f;
        const float kSolidAspectRatio = 0.2f;
        const int kSmoothingWindow = 1;

        [SerializeField] protected UVStyle m_uvStyle = UVStyle.Distance;
        [SerializeField] protected bool m_bOffsetInTexcoord1;

        public FlatGeometryBrush()
            : base(bCanBatch: true,
                upperBoundVertsPerKnot: kVertsInSolid,
                bDoubleSided: true)
        {
        }

        //
        // GeometryBrush API
        //

        protected override void InitBrush(BrushDescriptor desc,
                                          TrTransform localPointerXf)
        {
            base.InitBrush(desc, localPointerXf);
            SetDoubleSided(desc);
            m_geometry.Layout = GetVertexLayout(desc);
            m_sizes = new List<float>();
        }

        override public GeometryPool.VertexLayout GetVertexLayout(BrushDescriptor desc)
        {
            return new GeometryPool.VertexLayout
            {
                bUseColors = true,
                bUseNormals = true,
                bUseTangents = true,
                uv0Size = 2,
                uv0Semantic = GeometryPool.Semantic.XyIsUv,
                uv1Size = m_bOffsetInTexcoord1 ? 3 : 0,
                uv1Semantic = m_bOffsetInTexcoord1
                    ? GeometryPool.Semantic.Vector
                    : GeometryPool.Semantic.Unspecified
            };
        }

        override public float GetSpawnInterval(float pressure01)
        {
            return kSolidMinLengthMeters_PS * POINTER_TO_LOCAL * App.METERS_TO_UNITS +
                (PressuredSize(pressure01) * kSolidAspectRatio);
        }

        override protected void ControlPointsChanged(int iKnot0)
        {
            // Frames knots, determines how much geometry each knot should get
            OnChanged_FrameKnots(iKnot0);
            ResizeGeometry();

            // Updating a control point affects geometry generated by previous knot
            // (if there is any). The HasGeometry check is not a micro-optimization:
            // it also keeps us from backing up past knot 0.
            {
                // Vertices within a window of iKnot0 are affected. Walk backwards to find the first affected
                // index. Also, go back one extra because the tangent at a knot depends on its previous knot.
                int iKnotStart = iKnot0;
                for (int iKnotStartTest = iKnotStart - 1; iKnotStartTest >= iKnot0 - kSmoothingWindow - 1; iKnotStartTest--)
                {
                    if (iKnotStartTest < 0 || !m_knots[iKnotStartTest].HasGeometry)
                    {
                        break;
                    }
                    iKnotStart = iKnotStartTest;
                }
                OnChanged_MakeVertsAndNormals(iKnotStart);
            }

            // Distance between knots isn't changed by the geometry, so there's not (yet?)
            // any reason to back up when creating them.
            // TODO: Distances are now changed slightly because of the smoothing. Look into
            // whether it's significant enough to warrant rewriting the UV calculations.
            if (m_uvStyle == UVStyle.Stretch)
            {
                // Fix up the terminating UVs of the previous segment if we've just started a new segment.
                int start = (iKnot0 > 2 && !m_knots[iKnot0 - 1].HasGeometry) ? iKnot0 - 2 : iKnot0;
                OnChanged_StretchUVs(start);
            }
            else
            {
                OnChanged_DistanceUVs(iKnot0);
            }

            {
                int start = (m_knots[iKnot0 - 1].HasGeometry) ? iKnot0 - 1 : iKnot0;
                OnChanged_Tangents(start);
            }
        }

        /// Goes backwards looking for a break - if it doesn't have to look far, it deletes all knots
        /// after the break, resizes the geometry and then finalizes the brush.
        /// It's a bit weird that the control points will still exists for this invisible bit.
        private void TrimShortStrokeAfterBreak()
        {
            if (!m_bM11Compatibility)
            {
                int lastIndex = m_knots.Count - 1;
                while (lastIndex > 0 && m_knots[lastIndex].HasGeometry)
                {
                    lastIndex--;
                }
                if (lastIndex > 1 && (m_knots.Count - lastIndex) < kMinimumKnotsAfterBreak)
                {
                    m_knots.RemoveRange(lastIndex + 1, (m_knots.Count - lastIndex) - 1);
                    ResizeGeometry();
                }
            }
        }

        override public void FinalizeSolitaryBrush()
        {
            TrimShortStrokeAfterBreak();
            base.FinalizeSolitaryBrush();
        }

        override public BatchSubset FinalizeBatchedBrush()
        {
            TrimShortStrokeAfterBreak();
            return base.FinalizeBatchedBrush();
        }

        // Fills in any knot data needed for geometry generation.
        // - fill in length, nRight, nSurface, iVert, iTri
        // - calculate strip-break points
        void OnChanged_FrameKnots(int iKnot0)
        {
            float minMove = kMinMoveLengthMeters_PS * App.METERS_TO_UNITS * POINTER_TO_LOCAL;

            // A knot's tangent is affected by the previous and next knot positions so try to start at
            // iKnot0-1 if possible.
            int iKnotStart = m_knots[iKnot0 - 1].HasGeometry ? iKnot0 - 1 : iKnot0;
            Knot prev = m_knots[iKnotStart - 1];
            for (int iKnot = iKnotStart; iKnot < m_knots.Count; ++iKnot)
            {
                Knot cur = m_knots[iKnot];

                bool shouldBreak = false;

                Vector3 vMove = cur.point.m_Pos - prev.point.m_Pos;
                cur.length = vMove.magnitude;

                // Rather than use unstable math, just bail and don't change the geometry.
                if (cur.length < minMove)
                {
                    shouldBreak = true;
                }

                // invariant: nSurface = nMove x nRight
                // If single-sided, always point the frontside towards the brush. Causes twisting.
                Vector3 nTangent = vMove / cur.length;
                if (!m_bM11Compatibility && iKnot < m_knots.Count - 1)
                {
                    // Calculate a smoother tangent.
                    // TODO: Look into whether we should do something more accurate with the tangent
                    // like a cubic spline.
                    Knot next = m_knots[iKnot + 1];
                    nTangent = (next.point.m_Pos - prev.point.m_Pos).normalized;

                    if (Vector3.Dot(vMove, next.point.m_Pos - cur.point.m_Pos) < 0)
                    {
                        shouldBreak = true;
                    }
                }
                Vector3 vPreferredRight = m_Desc.m_BackIsInvisible
                    ? Vector3.Cross(cur.point.m_Orient * Vector3.forward, nTangent)
                    : prev.nRight;
                ComputeSurfaceFrameNew(
                    vPreferredRight, nTangent, cur.point.m_Orient,
                    out cur.nRight, out cur.nSurface);

                // More break checking; replicates previous logic
                if (m_bM11Compatibility && prev.HasGeometry)
                {
                    float fWidthHeightRatio = cur.length / PressuredSize(cur.smoothedPressure);
                    float fBreakAngle = Mathf.Atan(fWidthHeightRatio) * Mathf.Rad2Deg * kBreakAngleScalar;
                    Vector3 vPrevMove = prev.point.m_Pos - m_knots[iKnot - 2].point.m_Pos;
                    if (Vector3.Angle(vPrevMove, vMove) > fBreakAngle)
                    {
                        shouldBreak = true;
                    }
                }

                if (shouldBreak)
                {
                    // End of strip
                    cur.iTri = prev.iTri + prev.nTri;
                    cur.nTri = 0;
                    cur.iVert = (ushort)(prev.iVert + prev.nVert);
                    cur.nVert = 0;
                    cur.nRight = cur.nSurface = Vector3.zero;
                }
                else
                {
                    // Beginning or middle of strip
                    cur.iTri = prev.iTri + prev.nTri;
                    cur.iVert = (ushort)(prev.iVert + prev.nVert);
                    if (prev.HasGeometry)
                    {
                        // Middle of strip -- back up and share geometry
                        cur.iTri -= kSharedTrisInSolidPair * NS;
                        cur.iVert -= (ushort)(kSharedVertsInSolidPair * NS);
                    }
                    cur.nTri = (ushort)(kTrisInSolid * NS);
                    cur.nVert = (ushort)(kVertsInSolid * NS);
                }

                m_knots[iKnot] = cur;
                prev = cur;
            }
        }

        void OnChanged_MakeVertsAndNormals(int iKnot0)
        {
            // m_sizes is an array of sizes that corresponds to m_knots. If the array is too small, fill it
            // up with any size. It will be overwritten later.
            while (m_sizes.Count < m_knots.Count)
            {
                m_sizes.Add(0);
            }

            // Invariant: there is a previous knot.
            Knot prev = m_knots[iKnot0 - 1];
            float sizePrev = m_sizes[iKnot0 - 1];

            // Run through the knots to create the array of sizes.
            for (int iKnot = iKnot0; iKnot < m_knots.Count; ++iKnot)
            {
                // Invariant: all of prev's geometry (if any) is correct and up-to-date.
                // Thus, there is no need to modify anything shared with prev.
                Knot cur = m_knots[iKnot];

                if (cur.HasGeometry)
                {
                    SetTri(cur.iTri, cur.iVert, 0, BR, BL, FL);
                    SetTri(cur.iTri, cur.iVert, 1, BR, FL, FR);

                    if (!prev.HasGeometry)
                    {
                        // Can't use prev.nRight, prev.nSurface; they're invalid if no geometry
                        float size = PressuredSize(prev.smoothedPressure);
                        float alpha = PressuredOpacity(prev.smoothedPressure);
                        Vector3 halfRight = cur.nRight * (size / 2);
                        SetVert(cur.iVert, BR, prev.point.m_Pos + halfRight, cur.nSurface, m_Color, alpha);
                        SetVert(cur.iVert, BL, prev.point.m_Pos - halfRight, cur.nSurface, m_Color, alpha);
                        if (m_bOffsetInTexcoord1)
                        {
                            SetUv1(cur.iVert, BR, halfRight);
                            SetUv1(cur.iVert, BL, -halfRight);
                        }
                        sizePrev = m_sizes[iKnot - 1] = size;
                    }

                    if (m_bM11Compatibility)
                    {
                        float size = PressuredSize(cur.smoothedPressure);
                        float alpha = PressuredOpacity(cur.smoothedPressure);
                        Vector3 halfRight = cur.nRight * (size / 2);
                        SetVert(cur.iVert, FR, cur.point.m_Pos + halfRight, cur.nSurface, m_Color, alpha);
                        SetVert(cur.iVert, FL, cur.point.m_Pos - halfRight, cur.nSurface, m_Color, alpha);
                        if (m_bOffsetInTexcoord1)
                        {
                            SetUv1(cur.iVert, FR, halfRight);
                            SetUv1(cur.iVert, FL, -halfRight);
                        }
                    }
                    else
                    {
                        float size = PressuredSize(cur.smoothedPressure);

                        // Shrink the widths so that they don't self-intersect.
                        Vector3 prevForward = Vector3.Cross(prev.nRight, prev.nSurface);
                        float dotRight = Vector3.Dot(prevForward, cur.point.m_Pos + .5f * size * cur.nRight - prev.point.m_Pos);
                        float dotLeft = Vector3.Dot(prevForward, cur.point.m_Pos - .5f * size * cur.nRight - prev.point.m_Pos);
                        if ((dotLeft < 0 && dotRight > 0) || (dotLeft > 0 && dotRight < 0))
                        {
                            // Shrink the brush so that it doesn't self intersect.
                            Vector3 vEndPointLeft = prev.point.m_Pos - 0.5f * sizePrev * prev.nRight;
                            Vector3 vEndPointRight = prev.point.m_Pos + 0.5f * sizePrev * prev.nRight;
                            Vector3 clippedRight;
                            if (dotLeft < 0)
                            {
                                // Turning towards left side.
                                clippedRight = cur.point.m_Pos - vEndPointLeft;
                            }
                            else
                            {
                                // Turning towards right side.
                                clippedRight = vEndPointRight - cur.point.m_Pos;
                            }
                            size = clippedRight.magnitude;
                        }

                        // Make sure that size doesn't grow too quickly.
                        float moveLength = Vector3.Distance(cur.point.m_Pos, prev.point.m_Pos);
                        if (size > sizePrev + moveLength)
                        {
                            size = sizePrev + moveLength;
                        }

                        sizePrev = m_sizes[iKnot] = size;
                    }
                }

                prev = cur;
            }

            if (!m_bM11Compatibility)
            {
                // Run through the knots again to set the vertices based on the original knots and the
                // resulting size array.
                Vector3 halfRightPrev = m_knots[iKnot0 - 1].nRight * m_sizes[iKnot0 - 1] / 2;
                Vector3 knotPointPrev = m_knots[iKnot0 - 1].point.m_Pos;
                Vector3 halfRightCur = m_knots[iKnot0].nRight * m_sizes[iKnot0] / 2;
                Vector3 knotPointCur = m_knots[iKnot0].point.m_Pos;
                if (!m_knots[iKnot0 - 1].HasGeometry)
                {
                    // Can't use prev.nRight, prev.nSurface; they're invalid if no geometry
                    halfRightPrev = halfRightCur;
                }
                for (int iKnot = iKnot0; iKnot < m_knots.Count; ++iKnot)
                {
                    // Invariant: all of prev's geometry (if any) is correct and up-to-date.
                    // Thus, there is no need to modify anything shared with prev.
                    Knot cur = m_knots[iKnot];
                    int iNext = iKnot < m_knots.Count - 1 ? iKnot + 1 : iKnot;
                    Vector3 halfRightNext = m_knots[iNext].nRight * m_sizes[iNext] / 2;
                    Vector3 knotPointNext = m_knots[iNext].point.m_Pos;

                    if (cur.HasGeometry)
                    {
                        float alpha = PressuredOpacity(cur.smoothedPressure);
                        Vector3 surface = cur.nSurface;
                        Vector3 knotPoint = 0.3f * knotPointPrev + 0.4f * knotPointCur + 0.3f * knotPointNext;
                        Vector3 halfRight = 0.3f * halfRightPrev + 0.4f * halfRightCur + 0.3f * halfRightNext;
                        // The last nRight vector is very jittery, so don't do interpolation for the last couple
                        // knots.
                        if (!m_knots[iNext].HasGeometry /* || iKnot >= m_knots.Count - 2*/)
                        {
                            halfRight = halfRightCur;
                        }
                        SetVert(cur.iVert, FL, knotPoint - halfRight, surface, m_Color, alpha);
                        SetVert(cur.iVert, FR, knotPoint + halfRight, surface, m_Color, alpha);
                        if (m_bOffsetInTexcoord1)
                        {
                            SetUv1(cur.iVert, FR, halfRightCur);
                            SetUv1(cur.iVert, FL, -halfRightCur);
                        }
                    }

                    halfRightPrev = halfRightCur;
                    halfRightCur = halfRightNext;
                    knotPointPrev = knotPointCur;
                    knotPointCur = knotPointNext;
                }
            }
        }

        void OnChanged_StretchUVs(int iChangedKnot)
        {
            // Back up knot to the start of the segment
            // Invariant: knot 0 never has geometry
            int iKnot0 = iChangedKnot;
            while (m_knots[iKnot0 - 1].HasGeometry)
            {
                iKnot0 -= 1;
            }

            while (iKnot0 < m_knots.Count)
            {
                // Update UVs for a single segment

                // Invariant: If iKnot0-1 exists, it has no geometry
                // IOW, iKnot0 is not the middle of a strip.
                // (We don't assume that it does have geometry, though)

                // Find length and end knot of segment
                float totalLength = 0;
                int iKnot1 = iKnot0;
                for (; iKnot1 < m_knots.Count; ++iKnot1)
                {
                    Knot cur = m_knots[iKnot1];
                    if (!cur.HasGeometry)
                    {
                        break;
                    }
                    totalLength += cur.length;
                }

                // Invariant: all knots in [iKnot0, iKnot1) have geometry
                // If iKnot1 exists, it does not have geometry

                float v0, v1;
                {
                    float random01 = m_rng.In01(m_knots[iKnot0].iVert - 1);
                    int numV = m_Desc.m_TextureAtlasV;
                    int iAtlas = (int)(random01 * 3331) % numV;
                    v0 = (iAtlas) / (float)numV;
                    v1 = (iAtlas + 1) / (float)numV;
                }

                float distance = 0;
                for (int iKnot = iKnot0; iKnot < iKnot1; ++iKnot)
                {
                    Knot cur = m_knots[iKnot];
                    distance += cur.length;
                    float u = distance / totalLength;
                    if (iKnot == iKnot0)
                    {
                        SetUv0(cur.iVert, BL, new Vector2(0, v0));
                        SetUv0(cur.iVert, BR, new Vector2(0, v1));
                    }
                    SetUv0(cur.iVert, FL, new Vector2(u, v0));
                    SetUv0(cur.iVert, FR, new Vector2(u, v1));
                }

                iKnot0 = iKnot1 + 1;
            }
        }

        void OnChanged_DistanceUVs(int iKnot0)
        {
            Knot prev = m_knots[iKnot0 - 1];
            for (int iKnot = iKnot0; iKnot < m_knots.Count; ++iKnot)
            {
                Knot cur = m_knots[iKnot];

                if (cur.HasGeometry)
                {
                    // Textures are laid out so u goes along the strip,
                    // and v goes across the strip (from left to right)
                    float u0, v0, v1;
                    if (prev.HasGeometry)
                    {
                        // Middle of strip
                        Vector2 u0v0 = m_geometry.m_Texcoord0.v2[prev.iVert + FL * NS];
                        u0 = u0v0.x;
                        v0 = u0v0.y;
                        v1 = m_geometry.m_Texcoord0.v2[prev.iVert + FR * NS].y;
                    }
                    else
                    {
                        // Beginning of strip. Guaranteed no sharing, so previous
                        // vert should have a stable position.
                        float random01 = m_rng.In01(cur.iVert - 1);
                        u0 = random01;
                        int numV = m_Desc.m_TextureAtlasV;
                        int iAtlas = (int)(random01 * 3331) % numV;
                        v0 = (iAtlas) / (float)numV;
                        v1 = (iAtlas + 1) / (float)numV;

                        SetUv0(cur.iVert, BL, new Vector2(u0, v0));
                        SetUv0(cur.iVert, BR, new Vector2(u0, v1));
                    }

                    float length = m_knots[iKnot].length;
                    float size = PressuredSize(m_knots[iKnot].smoothedPressure);
                    float u1 = u0 + m_Desc.m_TileRate * (length / size);
                    SetUv0(cur.iVert, FL, new Vector2(u1, v0));
                    SetUv0(cur.iVert, FR, new Vector2(u1, v1));
                }

                prev = cur;
            }
        }

        void OnChanged_Tangents(int iKnot0)
        {
            Knot prev = m_knots[iKnot0 - 1];
            for (int iKnot = iKnot0; iKnot < m_knots.Count; ++iKnot)
            {
                Knot cur = m_knots[iKnot];
                if (cur.HasGeometry)
                {
                    // ComputeST is static API; takes vert indices instead of vert-pair indices.
                    Vector3 vS_BR_BL_FL, vS_BR_FL_FR, unused;
                    ComputeST(m_geometry.m_Vertices, m_geometry.m_Texcoord0.v2, cur.iVert, BR * NS, BL * NS, FL * NS,
                        out vS_BR_BL_FL, out unused);
                    ComputeST(m_geometry.m_Vertices, m_geometry.m_Texcoord0.v2, cur.iVert, BR * NS, FL * NS, FR * NS,
                        out vS_BR_FL_FR, out unused);

                    if (!prev.HasGeometry)
                    {
                        SetTangent(cur.iVert, BL, vS_BR_BL_FL);
                        SetTangent(cur.iVert, BR, vS_BR_BL_FL + vS_BR_FL_FR);
                    }

                    SetTangent(cur.iVert, FL, vS_BR_BL_FL + vS_BR_FL_FR);
                    SetTangent(cur.iVert, FR, vS_BR_FL_FR);
                }

                prev = cur;
            }
        }
    }
} // namespace TiltBrush
